package llm

import (
	"encoding/json"
	"strings"
	"time"

	"github.com/golang-jwt/jwt"
)

/* 智谱清言 */

const glm_api_key = "."

var glm_id_secret = strings.Split(glm_api_key, ".")

var glm_id = glm_id_secret[0]
var glm_secret = []byte(glm_id_secret[1])

const glm4_url = "https://open.bigmodel.cn/api/paas/v4/chat/completions"

/** glm-4 model */
type Glm4 struct {
	// token.expire is milli second timestamp
	Token
}

func (m *Glm4) RefreshToken() error {

	if m.expire > time.Now().UnixMilli()+5000 {
		return nil
	}

	now := time.Now()
	exp := time.Now().Add(time.Second * time.Duration(86400)).UnixMilli()

	payload := jwt.MapClaims{
		"api_key":   glm_id,
		"exp":       exp,
		"timestamp": now.UnixMilli(),
	}

	token := jwt.NewWithClaims(jwt.SigningMethodHS256, payload)
	token.Header["alg"] = "HS256"
	token.Header["sign_type"] = "SIGN"

	s, err := token.SignedString(glm_secret)

	if err != nil {
		return err
	}

	m.token = s
	m.expire = exp

	return nil
}

func (m *Glm4) CallApi(system_prompt, user_prompt string) (string, error) {

	rsp, err := client.R().
		SetHeader("Authorization", m.token).
		SetHeader("Content-Type", "application/json").
		SetBody(generateGlmReq(system_prompt, user_prompt)).
		Post(glm4_url)

	if err != nil {
		return "", err
	}

	var glm4Rsp GlmRsp
	err = json.Unmarshal(rsp.Body(), &glm4Rsp)
	if err != nil {
		return "", err
	}

	return glm4Rsp.Choices[0].Message.Content, nil
}

func generateGlmReq(system_prompt, user_prompt string) map[string]interface{} {
	return map[string]interface{}{
		"model": "glm-4",
		"messages": []map[string]string{
			{
				"role":    "system",
				"content": system_prompt,
			},
			{
				"role":    "user",
				"content": user_prompt,
			},
		},
	}
}

type GlmRsp struct {
	ID      string   `json:"id"`
	Created int64    `json:"created"`
	Model   string   `json:"model"`
	Choices []Choice `json:"choices"`
	Usage   Usage    `json:"usage"`
}

type Choice struct {
	Index        int     `json:"index"`
	FinishReason string  `json:"finish_reason"`
	Message      Message `json:"message"`
}

type Usage struct {
	PromptTokens     int `json:"prompt_tokens"`
	CompletionTokens int `json:"completion_tokens"`
	TotalTokens      int `json:"total_tokens"`
}

type Message struct {
	Role      string     `json:"role"`
	Content   string     `json:"content"`
	ToolCalls []ToolCall `json:"tool_calls"`
}

type ToolCall struct {
	Function `json:"function"`
	ID       string `json:"id"`
	Type     string `json:"type"`
}

type Function struct {
	Name   string      `json:"name"`
	Params interface{} `json:"arguments"`
}
